import sys
import av
import time
import datetime
import traceback
import numpy as np
from pybaseutils import image_utils
import cv2


class RTSPCamera(object):
    def __init__(self, name, url):
        self.name = name  # 视频名称
        self.url = url  # 视频URL或路径
        self.inp_av = None  # 输入视频流
        self.out_av = None  # 输出视频流
        self.video_stream = None  # 添加视频流引用
        self.audio_stream = None  # 添加音频流引用
        self.is_recording = False
        self.w = None
        self.h = None
        self.fps = None

    def connect(self):
        try:
            print("open url: " + self.url)
            # 设置连接超时和读取超时为N秒
            options = {
                'stimeout': '10000000',  # 设置socket超时为5秒 (单位：微秒)
                'timeout': '10000000'  # 设置读取超时为5秒 (单位：微秒)
            }
            self.inp_av = av.open(self.url, options=options)
            # Get video stream and its properties
            video_stream = next((s for s in self.inp_av.streams if s.type == 'video'), None)
            self.w = video_stream.width
            self.h = video_stream.height
            self.fps = video_stream.guessed_rate
            print("video info：({}x{}),fps={}".format(self.w, self.h, self.fps))
            return True
        except Exception as e:
            traceback.print_exc()
            print(f"连接错误: {str(e)}")
            return False

    def disconnect(self):
        if self.inp_av:
            self.inp_av.close()
            self.inp_av = None

    def start_recording(self):
        if not self.is_recording:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{self.name}_{timestamp}.mp4"
            self.out_av = av.open(filename, mode='w')

            # Get original video stream parameters
            inp_stream = next((s for s in self.inp_av.streams if s.type == 'video'), None)
            if inp_stream:
                # 添加视频流，使用原始分辨率
                self.video_stream = self.out_av.add_stream('h264', rate=self.fps)
                self.video_stream.width = self.w
                self.video_stream.height = self.h
                self.video_stream.pix_fmt = 'yuv420p'

            # 添加音频流（如果输入源有音频）
            if any(s.type == 'audio' for s in self.inp_av.streams):
                self.audio_stream = self.out_av.add_stream('aac', rate=48000)

            self.is_recording = True
            print("save video {}".format(filename))

    def stop_recording(self):
        if self.is_recording:
            self.out_av.close()
            self.out_av = None
            self.is_recording = False

    def start_display(self):
        if not self.connect():
            print("Failed to connect to camera")
            return
        for frame in self.inp_av.decode():
            if isinstance(frame, av.VideoFrame):  # Check if it's a video frame
                frame = frame.to_ndarray(format='bgr24')
                self.show_image(self.name, frame)

    def show_image(self, title, image, delay=10):
        if image is None: return image
        cv2.namedWindow(title, flags=cv2.WINDOW_NORMAL)
        cv2.imshow(title, image)
        key = cv2.waitKey(delay)
        if key & 0xFF == ord('q'):
            exit()
        return image


if __name__ == '__main__':
    # url = "rtsp://admin:C2332416@192.168.2.35"
    url = "../test-video.mp4"
    c = RTSPCamera(name="demo", url=url)
    c.start_display()
